import dlib,os,glob,time
import  cv2
import numpy as np
import csv
import pandas as pd
import math
from scipy.spatial.distance import pdist
import matplotlib.pyplot as plt
import json

class MultiFaceReco():

    #计算128D描述符的欧式距离
    def compute_dst(self,feature_1,feature_2):
        feature_1 = np.array(feature_1)
        feature_2 = np.array(feature_2)
        #dist1 = pdist(np.vstack([feature_1,feature_2]),'cosine')
        dist = np.linalg.norm(feature_1 - feature_2)
        return dist

    def get_facelibrary(self):
        featuremean_path = "./dataset/task1/featureMean/"
        head = [] # 没有表头的会导致将内容认为是表头
        for i in range(128):
            fe = "feature_" + str(i + 1)
            head.append(fe)
        face_path = featuremean_path+"feature_all.csv"
        face_feature=pd.read_csv(face_path,names=head)

        # print(face_feature.shape)

        face_feature_array=np.array(face_feature)
        #print(face_feature_array.shape)
        return face_feature_array

    def face_recognition(self,image):
        # 加载检测关键点定位和识别模型
        predictor_path = "./dataset/task1/model/shape_predictor_68_face_landmarks.dat"
        model_path = "./dataset/task1/model/dlib_face_recognition_resnet_model_v1.dat"

        detector = dlib.get_frontal_face_detector() # 框选人脸
        # 人脸关键点标记
        predictor= dlib.shape_predictor(predictor_path) # 返回68个特征点的位置
        # 生成面部识别器
        facerec = dlib.face_recognition_model_v1(model_path)
        
        face_mean = np.zeros(128)
        
        F,H,W,_ = image.shape
        #print(image.shape)
        boundary1 = W/3
        boundary2 = W/3*2
        count = np.zeros((3,20))
        finish = 0                    
        for i in range(F):
            if i % 6== 0: 
                finish += 1          # 每6帧截取一帧
                # 转为灰度图像处理
                gray = cv2.cvtColor(image[i,:,:,:], cv2.COLOR_BGR2GRAY)
                #print(np.array(gray).astype)
                #gray = sharpen(gray) 
                dets = detector(gray, 1)
                #print(np.array(dets))
                #print(dets)        # 检测帧图像中的人脸
                # 处理检测到的每一张人脸
                if len(dets)>0:
                    for Index,value in enumerate(dets): # 返回检测区域的矩形坐标
                        #for value in dets: 
                        #print(Index,value)
                        #获取面部关键点
                        shape = predictor(gray,value)
                        # 提取特征-图像中的68个关键点转换为128D面部描述符
                        face_descriptor = facerec.compute_face_descriptor(image[i,:,:,:], shape)
                        v = np.array(face_descriptor)
                        ID,index = self.Result(v,self.get_facelibrary())
                        if (value.center().x<boundary1):
                            Index_ = 0
                        if (boundary1<value.center().x<boundary2):
                            Index_ = 1
                        if (value.center().x>boundary2):
                            Index_ = 2
                        count[Index_,index] += 1 # 记录被分到每一个类别的帧数
                        #Index += 1
                #print("已识别该视频帧数的{:.2f}%".format(finish/(math.floor(F/6)+(F%6!=0))*100))
        Maxindex = np.argmax(count,axis=1)
        #print(count)
        #print(Maxindex)
                    #face_mean = face_mean+v
        #face_mean = face_mean/(math.floor(F/6))
        #return face_mean
        return Maxindex

    def Result(self,face_mean,face_feature_array): # 返回ID
        face_list = ["ID1","ID10","ID11","ID12","ID13","ID14","ID15","ID16","ID17","ID18","ID19","ID2","ID20","ID3","ID4","ID5","ID6","ID7","ID8","ID9"] # 根据遍历顺序得到
        #for i in range(1,21,1):
            #face_list.append("ID{}".format(i))
        distance_list = [] # 存储到各个平均特征的距离
        for i in range(20):
            distance_list.append(self.compute_dst(face_mean,face_feature_array[i]))

        result_ = distance_list.index(min(distance_list)) # 返回最小距离的索引
        result = face_list[result_]

        return result,result_

    def task3ID(self,video_path):
        face_list = ["ID1","ID10","ID11","ID12","ID13","ID14","ID15","ID16","ID17","ID18","ID19","ID2","ID20","ID3","ID4","ID5","ID6","ID7","ID8","ID9"] # 根据遍历顺序得到
        face_feature_array = self.get_facelibrary()
        # 测试1
        result_dict = {}
        for file_name in os.listdir(video_path):
            ## 读取MP4文件中的视频,可以用任意其他的读写库
            cap=cv2.VideoCapture(video_path+'/'+file_name)
            frames_num = int(cap.get(7))
            frames_H = int(cap.get(4))
            frames_W = int(cap.get(3))
            video_frames = np.zeros((frames_num,frames_H,frames_W,3))
            if cap.isOpened():
                ret, frame = cap.read()
            else:
                ret = False
            #frame_interval = 30  # 截取图片间隔帧的数量
            frame_count = 0 # 帧计数
            # 循环读取帧
            while ret:
                video_frames[frame_count,:,:,:] = frame
                ret, frame = cap.read()
                frame_count += 1
            cap.release()
            #print(video_frames.shape)
            video_frames = video_frames.astype(np.uint8)
            
            #face_mean = task_1_recognition.face_recognition(video_frames)
            #result = task_1_recognition.Result(face_mean,face_feature_array)
            Maxindex = self.face_recognition(video_frames)

            ## 返回一个ID
            result_dict[file_name]=[face_list[Maxindex[0]],face_list[Maxindex[1]],face_list[Maxindex[2]]]
            #print("{} is {}".format(file_name,face_list[Maxindex]))

        return result_dict

multifacereco = MultiFaceReco()
video_path = "./dataset/test_offline/task3"
dict = multifacereco.task3ID(video_path)
with open("./src/separationhandler/data/facedata.json","w") as f:
    json.dump(dict,f,indent=4, ensure_ascii=False)

